# 8个线性天线3层2次多项式拟合：12天线、sslDb=-30dB、扫描角度0度
# 将w做多项式处理
import numpy as np
from scipy.signal import windows as SSW # import chebwin
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from torch.optim.lr_scheduler import ReduceLROnPlateau
from apps.wfs.nlbf_config import NlbfConfig as NG

class NlbfModel(nn.Module):
    def __init__(self, N_base:int =8, rank:int =3, init_coefs:bool =False, avec_dim:int =1, layers:int = 1):
        super(NlbfModel, self).__init__()
        self.rank = rank
        self.N_base = N_base
        self.layers = layers
        # 定义一个N_base x rank的矩阵，每一行对应一个多项式的系数
        if init_coefs:
            self.load()
        else:
            for i in range(avec_dim):
                for j in range(self.layers):
                    self.register_parameter(f'c{j}_{i}', nn.Parameter(torch.randn(N_base, self.rank)))
    
    def forward(self, a, x=None, mode:int =1):
        '''
        a: ((1800, 8), complex128)
        x: ((8,), complex128)
        输出为：((1800,1), float64)
        '''
        outputs = None
        layer_output = None # 本层输出
        z = a # 上一层输出
        for layer_idx in range(self.layers):
            layer_output = None
            w = torch.conj(x).unsqueeze(1)
            w_pow = torch.stack([x**i for i in range(self.rank)], dim=1) # w_pow: ((8,3), complex128)
            w_poly = getattr(self, f'c{layer_idx}_{0}') * w_pow # w_poly: ((8,3), complex128)
            for idx in range(self.rank):
                layer_output = z*w_poly[:, idx] # layer_output: ((1800,8), complex128)
            else:
                layer_output += z*w_poly[:, idx]
            z = layer_output # z: ((1800,8), complex128)
        if mode == 1:
            outputs_raw = torch.matmul(layer_output, torch.conj(x).unsqueeze(1)) # outputs: ((1800,1), float64)
        else:
            outputs_raw = torch.matmul(torch.conj(x).unsqueeze(1).T, layer_output.T) # outputs: ((1800,1), float64)
        outputs = torch.abs(outputs_raw) / torch.max(torch.abs(outputs_raw))
        return outputs

    def save(self, pt_fn='./work/wfs/nlbf.pt') -> None:
        torch.save(self.state_dict(), pt_fn)

    def load(self, pt_fn='./work/wfs/nlbf.pt') -> None:
        self.load_state_dict(torch.load(pt_fn))

    def get_weights(self) -> torch.tensor:
        return self.coefficients[0]     